
package com.trydofor.id.mina;

import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.util.ArrayList;
import java.util.List;

import org.apache.mina.core.filterchain.IoFilterAdapter;
import org.apache.mina.core.session.IdleStatus;
import org.apache.mina.core.session.IoSession;
import org.apache.mina.core.write.WriteRequest;
import org.apache.mina.filter.firewall.Subnet;
import org.slf4j.LoggerFactory;

import com.trydofor.aloe.exception.Exceptions;
import com.trydofor.id.logs.Eventable;
import com.trydofor.id.logs.Eventer;
import com.trydofor.id.logs.LogLevel;

/**
 * A {@link org.apache.mina.core.filterchain.IoFilter} which blocks connections Not from whitelisted.
 * 
 * @author : Shi Rongjiu (www.trydofor.com)
 */
public class WhitelistFilter extends IoFilterAdapter implements Eventable {
    
    /**
     * the events
     */
    public enum Event implements Eventer.Event {
        /**
         * block a session that is Not in the white list.
         */
        BLOCK_SESSION
    }
    
    private final Eventer      eventer   = new Eventer(LoggerFactory.getLogger(WhitelistFilter.class));
    private final List<Subnet> whitelist = new ArrayList<Subnet>();
    
    /**
     * constructor
     */
    public WhitelistFilter() {
        eventer.register(Event.BLOCK_SESSION, LogLevel.WARN);
    }
    
    /**
     * {@inheritDoc}
     */
    public Eventer getEventer() {
        return eventer;
    }
    
    /**
     * allow the specified ipv4 address.
     * 
     * @param address the ipv4 address
     * @throws NullPointerException if address is null
     */
    public void allow(InetAddress address) {
        Exceptions.checkNullPointer(address);
        allow(new Subnet(address, 32));
    }
    
    /**
     * allow the specified ipv4 subnet.
     * 
     * @param subnet the ipv4 subnet
     * @throws NullPointerException if address is null
     */
    public void allow(Subnet subnet) {
        Exceptions.checkNullPointer(subnet);
        whitelist.add(subnet);
    }
    
    /**
     * block the specified ipv4 address.
     * 
     * @param address the ipv4 address
     * @throws NullPointerException if address is null
     */
    public void remove(InetAddress address) {
        Exceptions.checkNullPointer(address);
        remove(new Subnet(address, 32));
    }
    
    /**
     * block the specified ipv4 subnet.
     * 
     * @param subnet the ipv4 subnet
     * @throws NullPointerException if address is null
     */
    public void remove(Subnet subnet) {
        Exceptions.checkNullPointer(subnet);
        whitelist.remove(subnet);
    }
    
    /**
     * copy the white list
     * 
     * @return the copy
     */
    public List<Subnet> copyWhiteList() {
        List<Subnet> copy = new ArrayList<Subnet>();
        copy.addAll(whitelist);
        return copy;
    }
    
    /**
     * {@inheritDoc}
     */
    public void sessionCreated(NextFilter nextFilter, IoSession session) {
        if (isNextElseClosing(session)) {
            nextFilter.sessionCreated(session);
        }
    }
    
    /**
     * {@inheritDoc}
     */
    public void sessionOpened(NextFilter nextFilter, IoSession session) throws Exception {
        if (isNextElseClosing(session)) {
            nextFilter.sessionOpened(session);
        }
    }
    
    /**
     * {@inheritDoc}
     */
    public void sessionClosed(NextFilter nextFilter, IoSession session) throws Exception {
        if (isNextElseClosing(session)) {
            nextFilter.sessionClosed(session);
        }
    }
    
    /**
     * {@inheritDoc}
     */
    public void sessionIdle(NextFilter nextFilter, IoSession session, IdleStatus status) throws Exception {
        if (isNextElseClosing(session)) {
            nextFilter.sessionIdle(session, status);
        }
    }
    
    /**
     * {@inheritDoc}
     */
    public void messageReceived(NextFilter nextFilter, IoSession session, Object message) {
        if (isNextElseClosing(session)) {
            nextFilter.messageReceived(session, message);
        }
    }
    
    /**
     * {@inheritDoc}
     */
    public void messageSent(NextFilter nextFilter, IoSession session, WriteRequest writeRequest)
            throws Exception {
        if (isNextElseClosing(session)) {
            nextFilter.messageSent(session, writeRequest);
        }
    }
    
    // ----------------------------
    
    private boolean isNextElseClosing(IoSession session) {
        
        SocketAddress remoteAddress = session.getRemoteAddress();
        
        if (remoteAddress instanceof InetSocketAddress) {
            InetAddress address = ((InetSocketAddress) remoteAddress).getAddress();
            for (Subnet subnet : whitelist) {
                if (subnet.inSubnet(address)) return true;
            }
        }
        
        // block it
        eventer.log(Event.BLOCK_SESSION, "Remote address [{}] NOT in  the whitelist, closing.", remoteAddress);
        
        session.close(true);
        return false;
    }
}
